﻿
#include "armature_node.h"
#include <assimp/scene.h> // 包含Assimp的场景结构体
#include <game/animation/avatar.h>
#include <graphics/imath.h>
#include <kfile.h>
#include <klog.h>

using namespace DirectX;

namespace ifire::game {

static const char* ASSIMP_PRE_NODE = "$AssimpFbx$";

ArmatureNode::ArmatureNode(ArmatureTree* _owner, int _index,
    ArmatureNode* parent_node, int child_count, ikit::sstr _name,
    FMat _transformation)
    : owner(_owner)
    , index(_index)
    , parent(parent_node)
    , name(_name)
    , children(child_count)
    , transformation(_transformation)
    , bone_id(-1)
    , avatar_id(AVATAR_NONE_ID)
    , tpose_relative(ID_MAT)
    , tpose_mat(ID_MAT)
    , tpose_length(0.0f) {}

bool ArmatureNode::IsRootMotionNode() const {
  return index == owner->root_motion_node_->index;
}

ArmatureTree::ArmatureTree(const aiNode* root_node, bool convert_units)
    : convert_units_(convert_units), is_set_avatar(false), is_calc_pose(false) {
  // 默认骨骼根数，其实应该有更聪明的办法获得
  nodes_.reserve(256);

  // 通过CreateFrom，可以直接创建所有的项目
  root_node_ = CreateFrom(nullptr, root_node);

  // RootMotion默认和项目一样
  root_motion_node_ = root_node_;
}

ArmatureTree::ArmatureTree(const ikit::XmlNode& root_node) {
  // 载入每一个子项目
  root_node.Foreach("Bone", [this](ikit::XmlNode& n, int idx) {
    int index = static_cast<int>(nodes_.size());
    // 这里绝对能够获得ParentNode，因为逻辑关系保证ParentNode肯定是在前面载入
    auto parent_node = Find(n("Parent"));
    XMFLOAT4X4 m;
    sscanf(n.Value(), "%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f", &m._11,
        &m._12, &m._13, &m._14, &m._21, &m._22, &m._23, &m._24, &m._31, &m._32,
        &m._33, &m._34, &m._41, &m._42, &m._43, &m._44);
    // 初始化项目
    auto item = new ArmatureNode(
        this, index, parent_node, 0, n("Name"), XMLoadFloat4x4(&m));
    nodes_.push_back(item);
  });

  BuildRelation();
}

ArmatureTree::ArmatureTree(ikit::CachedFile& f) {
  // 首先读取数量
  int count = f.readInt();

  char buffer[64];

  for(int i = 0; i < count; i++){
    // 按次序读入信息
    int index = f.readInt();
    int bone_id = f.readInt();
    ikit::sstr name = f.readString(buffer);
    ikit::sstr parent_name = f.readString(buffer);
    XMFLOAT4X4 mat;
    f.readBytes(&mat, sizeof(mat));

    // 初始化项目
    auto item = new ArmatureNode(
        this, index, Find(parent_name), 0, name, XMLoadFloat4x4(&mat));
    nodes_.push_back(item);

    // 配置其它的信息，这东西其实现在是-1，以后进行装配的，但为了数据完整性，还是写进去，
    // 因为比起顶点来说，这点多余的数据根本不算什么。
    item->bone_id = bone_id;
  }

  BuildRelation();
}

void ArmatureTree::BuildRelation() {
  // 用于获得子项目的参数
  auto get_parent_count = [this](const ArmatureNode* ptr) -> size_t {
    size_t count = 0;
    for (auto& i : nodes_) {
      if (i->parent == ptr) {
        count++;
      }
    }
    return count;
  };

  // 遍历项目，为其增加子项目
  for (auto& p : nodes_) {
    auto childen_count = get_parent_count(p);
    if (childen_count > 0) {
      p->children.reserve(childen_count);
      for (auto& i : nodes_) {
        if (i->parent == p) {
          p->children.push_back(i);
        }
      }
    }
  }

  // 默认第一位的项目，RootMotion默认和项目一样
  root_node_ = nodes_[0];
  root_motion_node_ = root_node_;
}

ArmatureTree::~ArmatureTree() {
  // 递归删除所有子节点
  for (ArmatureNode* i : nodes_) {
    delete i;
  }
  nodes_.clear(); // 清空子节点列表
}


void ArmatureTree::Save(ikit::CachedFile& f) {  
  f.writeInt(static_cast<int>(nodes_.size()));

  for (auto i : nodes_) {
    f.writeInt(i->index);
    f.writeInt(i->bone_id);
    f.writeString(i->name.c_str());
    if (i->parent){
      f.writeString(i->parent->name.c_str());
    }else{
      f.writeString(nullptr);
    }

    XMFLOAT4X4 mat;
    XMStoreFloat4x4(&mat, i->transformation);
    f.writeBytes(&mat, sizeof(mat));
  }
}

void ArmatureTree::SetRootMotionNode(ikit::sstr name) {
  auto n = Find(name);
  if (n) {
    root_motion_node_ = n;
  }
}

ArmatureNode* ArmatureTree::Find(ikit::sstr name) const {
  if (name.empty()) return nullptr;
  for (auto n : nodes_) {
    if (n->name == name) {
      return n;
    }
  }
  return nullptr;
}

void ArmatureTree::PrintXml(const std::wstring& filename) const {
  ikit::file_printer f(filename);
  XMFLOAT3 p, r;

  for (auto i : nodes_) {
    Trans3 t(i->transformation);
    p = t.GetPosition();
    r = t.GetRotationEuler();
    f.print(
        "<Bone Name='%s' Pos='%.5f,%.5f,%.5f' Rot='%.5f,%.5f,%.5f'></Bone>\n",
        i->name.c_str(), p.x, p.y, p.z, r.x, r.y, r.z);
  }
}

void ArmatureTree::PrintArmatureXml(const std::wstring& filename) const {
  ikit::file_printer f(filename);

  for (auto i : nodes_) {
    XMFLOAT4X4 m;
    XMStoreFloat4x4(&m, i->transformation);
    const char* parent = "";
    if (i->parent) {
      parent = i->parent->name.c_str();
    }
    f.print(
        "<Bone Name='%s' "
        "Parent='%s'>%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f,%f</Bone>\n",
        i->name.c_str(), parent, m._11, m._12, m._13, m._14, m._21, m._22,
        m._23, m._24, m._31, m._32, m._33, m._34, m._41, m._42, m._43, m._44);
  }
}

void ArmatureTree::PrintData(const char* file_name) const {
  ikit::file_printer f(file_name);
  XMFLOAT3 p, r;

  for (auto i : nodes_) {
    Trans3 t(i->transformation);
    p = t.GetPosition();
    r = t.GetRotationEuler();
    f.print("%s: P[%.3f,%.3f,%.3f] R[%.0f,%.0f,%.0f]\n", i->name.c_str(), p.x,
        p.y, p.z, r.x, r.y, r.z);
  }
}

void ArmatureTree::SetByAvatar(Avatar* avatar) {

  if (!is_set_avatar) {
    /*KLOG("\n=== SetByAvatar %s ===", avatar->name.c_str());*/

    for (auto e : nodes_) {
      e->avatar_id = avatar->FindId(e->name);
    }

    // 计算t-pose
    CalcPose(avatar);

    // 设置开始骨骼
    SetRootMotionNode(avatar->FindBone(AVATAR_ROOT_ID));

    is_set_avatar = true;

    /*KLOG("Width:%f, Height:%f, Long:%f", GetBoxWidth(), GetBoxHeight(),
        GetBoxLong());*/
  }
}

ArmatureNode* ArmatureTree::CreateFrom(
    ArmatureNode* parent_node, const aiNode* node) {
  // 获得当前的transform。根据是否需要变换单位，进行不同的处理。
  XMFLOAT4X4 t;
  if (convert_units_) {
    t = imath::ConvertUnitsMatrix(node->mTransformation);
  } else {
    t = imath::ConvertMatrix(node->mTransformation);
  }
  IMat transformation = XMLoadFloat4x4(&t);

  // 获得名称
  ikit::sstr name(node->mName.C_Str());

  // 检查名称中是否拥有"$AssimpFbx$"
  while (true) {
    // 如果没有，跳出循环
    if (name.index_of(ASSIMP_PRE_NODE) == -1 || node->mChildren == nullptr) {
      break;
    }

    // 得到下一个Node
    const aiNode* next = node->mChildren[0];
    // 与下一个Node进行计算
    XMFLOAT4X4 next_mat(imath::ConvertMatrix(next->mTransformation));
    auto next_m = XMLoadFloat4x4(&next_mat);
    // 进行乘法，注意是右乘，所以父级在右边
    // （其实这个关系不是很大，因为Assimp的处理，导致下一个Node只能是单位矩阵）
    transformation = XMMatrixMultiply(next_m, transformation);
    name = ikit::sstr(next->mName.C_Str());
    node = next;
  }

  // 返回项目
  // 根据Size得到索引
  int index = static_cast<int>(nodes_.size());
  auto item = new ArmatureNode(
      this, index, parent_node, node->mNumChildren, name, transformation);
  nodes_.push_back(item);

  // 如果有子项目，创建子项目
  for (unsigned int i = 0; i < node->mNumChildren; i++) {
    item->children[i] = CreateFrom(item, node->mChildren[i]);
  }

  // 返回项目
  return item;
}

void ArmatureTree::PrintHierarchy() const {
  KLOG("== PrintHierarchy ==");
  ikit::reset_tab();
  PrintHierarchy(root_node_);
}

void ArmatureTree::CalcPose(Avatar* avatar) {
  if (!is_calc_pose) {
    // KLOG("\n\n== %s T-POSE ==", avatar->name.c_str());

    // 计算t-pose
    CalcTPoseHierarchy(ID_MAT, root_node_, avatar, 0.0f);
    // 计算包围盒
    CalcBoxSize();

    is_calc_pose = true;
  }
}

void ArmatureTree::CalcBoxSize() {
  // 计算包围盒
  // 初始化为极值
  out_min_ = XMVectorReplicate(FLT_MAX);
  out_max_ = XMVectorReplicate(-FLT_MAX);

  // 遍历所有骨骼（假设骨骼原点在关节位置）
  for (size_t i = 0; i < nodes_.size(); ++i) {
    // 提取骨骼位置（矩阵第4行前3个分量）
    XMVECTOR bonePos = nodes_[i]->tpose_mat.r[3];

    // 更新包围盒
    out_min_ = XMVectorMin(out_min_, bonePos);
    out_max_ = XMVectorMax(out_max_, bonePos);
  }
  box_size_ = XMVectorSubtract(out_max_, out_min_);
}

void ArmatureTree::PrintHierarchy(const ArmatureNode* node) const {
  Trans3 t(node->transformation);

  ikit::print_tab();
  t.Print("[%d]%s B(%d)", node->index, node->name.c_str(), node->bone_id);

  ikit::add_tab();
  for (size_t i = 0; i < node->children.size(); i++) {
    PrintHierarchy(node->children[i]);
  }
  ikit::back_tab();
}

void ArmatureTree::CalcTPoseHierarchy(DirectX::FXMMATRIX parent,
    ArmatureNode* node, Avatar* avatar, float total_length) {
  // 获得Pose
  auto& pose = avatar->GetPoseBone(node->name);
  node->tpose_relative = pose.transform;

  // 记录当前层级
  DirectX::XMMATRIX current;

  // 如果Pose为空，则使用项目自已的Transform
  if (pose.bone.empty()) {
    auto node_trans = node->transformation;
    current = XMMatrixMultiply(node_trans, parent);
  } else {
    current = XMMatrixMultiply(pose.transform.GetMatrix(), parent);
    pose.final_mat = current;
  }

  node->tpose_mat = current;
  node->tpose_length = imath::CalculateBoneLength(current, parent);
  node->tpose_total_length = node->tpose_length + total_length;
  // KLOG("*%s*, Length: %f, Total: %f", node->name.c_str(), node->tpose_length,
  //     node->tpose_total_length);

  // 继续执行子项
  for (size_t i = 0; i < node->children.size(); i++) {
    CalcTPoseHierarchy(
        current, node->children[i], avatar, node->tpose_total_length);
  }
}

} // namespace ifire::game